package com.tarena.lbs.basic.controller;

import java.time.Duration;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

import com.alibaba.cloud.ai.dashscope.agent.DashScopeAgent;
import com.alibaba.cloud.ai.dashscope.agent.DashScopeAgentOptions;
import com.alibaba.cloud.ai.dashscope.api.DashScopeAgentApi;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

@RestController
@RequestMapping("/ai")
public class BailianAgentStreamController {
    private static final Logger logger = LoggerFactory.getLogger(BailianAgentStreamController.class);
    private DashScopeAgent agent;

    @Value("${spring.ai.dashscope.agent.options.app-id}")
    private String appId;

    public BailianAgentStreamController(DashScopeAgentApi dashscopeAgentApi) {
        this.agent = new DashScopeAgent(dashscopeAgentApi,
                DashScopeAgentOptions.builder()
                        .withSessionId("current_session_id")
                        .withIncrementalOutput(true)
                        .withHasThoughts(true)
                        .build());
    }

    @GetMapping(value="/bailian/agent/stream", produces=MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> stream(@RequestParam(value = "message") String message) {
        logger.info("收到流式请求，消息长度: {}", message.length());

        // 使用 Sinks 来更好地控制流
        Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();
        AtomicBoolean isCompleted = new AtomicBoolean(false);

        try {
            // 订阅 AI 流
            agent.stream(new Prompt(message, DashScopeAgentOptions.builder().withAppId(appId).build()))
                    .subscribe(
                            response -> {
                                try {
                                    if (response == null || response.getResult() == null) {
                                        logger.warn("chat response is null");
                                        sink.tryEmitNext("data: \n\n");
                                        return;
                                    }

                                    AssistantMessage app_output = response.getResult().getOutput();
                                    String content = app_output.getText();

                                    if (content != null && !content.trim().isEmpty()) {
                                        // 发送数据块
                                        String formattedContent = "data: " + content.replace("\n", "\\n") + "\n\n";
                                        sink.tryEmitNext(formattedContent);
                                        logger.debug("发送数据块，长度: {}", content.length());
                                    }
                                } catch (Exception e) {
                                    logger.error("处理响应时发生错误", e);
                                    if (!isCompleted.get()) {
                                        sink.tryEmitNext("data: [ERROR]处理响应时发生错误\n\n");
                                    }
                                }
                            },
                            error -> {
                                logger.error("流式响应发生错误", error);
                                if (!isCompleted.get()) {
                                    sink.tryEmitNext("data: [ERROR]服务暂时不可用，请稍后重试\n\n");
                                    sink.tryEmitNext("data: [DONE]\n\n");
                                    isCompleted.set(true);
                                }
                            },
                            () -> {
                                logger.info("AI流式响应正常完成");
                                if (!isCompleted.get()) {
                                    sink.tryEmitNext("data: [DONE]\n\n");
                                    isCompleted.set(true);
                                }
                            }
                    );

            // 返回 Flux 并添加超时和完成处理
            return sink.asFlux()
                    .timeout(Duration.ofSeconds(120))
                    .doOnCancel(() -> {
                        logger.info("客户端取消了流式请求");
                        if (!isCompleted.get()) {
                            sink.tryEmitNext("data: [DONE]\n\n");
                            isCompleted.set(true);
                        }
                    })
                    .doOnComplete(() -> {
                        logger.info("Flux流完成");
                        if (!isCompleted.get()) {
                            sink.tryEmitNext("data: [DONE]\n\n");
                            isCompleted.set(true);
                        }
                    });

        } catch (Exception e) {
            logger.error("创建流式响应失败", e);
            return Flux.just("data: [ERROR]创建请求失败\n\n", "data: [DONE]\n\n");
        }
    }
}